"""
知识库 API 路由

提供知识库相关的 REST API 端点
"""

from fastapi import APIRouter, HTTPException, File, UploadFile, Form
from pydantic import BaseModel
from typing import List, Dict, Optional
from module.knowledge_base import get_kb_manager, KnowledgeBaseManager
import json
import io
import logging
import traceback
import os

# 配置日志
logger = logging.getLogger(__name__)

# 尝试导入文档解析库
try:
    from docx import Document as DocxDocument
    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False
    logger.warning("python-docx 未安装，无法解析 .docx 文件")

try:
    import PyPDF2
    PDF_AVAILABLE = True
except ImportError:
    PDF_AVAILABLE = False
    logger.warning("PyPDF2 未安装，无法解析 .pdf 文件")

# 创建路由器
router = APIRouter(prefix="/api/knowledge-base", tags=["knowledge-base"])

# ==================== 文件解析函数 ====================

def parse_file_content(contents: bytes, filename: str, content_type: str) -> str:
    """
    根据文件类型解析文件内容
    
    Args:
        contents: 文件二进制内容
        filename: 文件名
        content_type: MIME 类型
        
    Returns:
        解析后的文本内容
    """
    # 获取文件扩展名
    ext = os.path.splitext(filename)[1].lower() if filename else ""
    
    logger.info(f"解析文件: filename={filename}, ext={ext}, content_type={content_type}")
    
    # 解析 .docx 文件
    if ext == ".docx" or content_type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
        if not DOCX_AVAILABLE:
            raise HTTPException(
                status_code=400, 
                detail="无法解析 .docx 文件，请安装 python-docx: pip install python-docx"
            )
        try:
            doc = DocxDocument(io.BytesIO(contents))
            paragraphs = []
            for para in doc.paragraphs:
                if para.text.strip():
                    paragraphs.append(para.text)
            # 提取表格内容
            for table in doc.tables:
                for row in table.rows:
                    row_text = " | ".join([cell.text.strip() for cell in row.cells])
                    if row_text.strip():
                        paragraphs.append(row_text)
            content = "\n".join(paragraphs)
            logger.info(f"成功解析 .docx 文件，提取了 {len(paragraphs)} 个段落")
            return content
        except Exception as e:
            logger.error(f"解析 .docx 文件失败: {str(e)}")
            logger.error(traceback.format_exc())
            raise HTTPException(status_code=400, detail=f"解析 .docx 文件失败: {str(e)}")
    
    # 解析 .pdf 文件
    elif ext == ".pdf" or content_type == "application/pdf":
        if not PDF_AVAILABLE:
            raise HTTPException(
                status_code=400,
                detail="无法解析 .pdf 文件，请安装 PyPDF2: pip install PyPDF2"
            )
        try:
            pdf_file = io.BytesIO(contents)
            pdf_reader = PyPDF2.PdfReader(pdf_file)
            text_parts = []
            for page_num, page in enumerate(pdf_reader.pages):
                try:
                    text = page.extract_text()
                    if text.strip():
                        text_parts.append(f"--- 第 {page_num + 1} 页 ---\n{text}")
                except Exception as e:
                    logger.warning(f"提取 PDF 第 {page_num + 1} 页失败: {str(e)}")
                    continue
            content = "\n\n".join(text_parts)
            logger.info(f"成功解析 .pdf 文件，提取了 {len(pdf_reader.pages)} 页")
            return content
        except Exception as e:
            logger.error(f"解析 .pdf 文件失败: {str(e)}")
            logger.error(traceback.format_exc())
            raise HTTPException(status_code=400, detail=f"解析 .pdf 文件失败: {str(e)}")
    
    # 解析文本文件 (.txt, .md 等)
    else:
        # 尝试多种编码
        encodings = ['utf-8', 'gbk', 'gb2312', 'latin-1']
        content_str = None
        
        for encoding in encodings:
            try:
                content_str = contents.decode(encoding)
                logger.info(f"文本文件解码成功，使用编码: {encoding}")
                break
            except UnicodeDecodeError:
                continue
        
        if content_str is None:
            # 如果所有编码都失败，使用 utf-8 并替换错误字符
            logger.warning("所有编码尝试失败，使用 utf-8 并替换错误字符")
            content_str = contents.decode("utf-8", errors='replace')
        
        return content_str

# ==================== 数据模型 ====================

class DocumentInput(BaseModel):
    """文档输入模型"""
    content: str
    source: str
    doc_type: str
    metadata: Optional[Dict] = None


class DocumentBatchInput(BaseModel):
    """批量文档输入"""
    documents: List[DocumentInput]


class SearchQuery(BaseModel):
    """搜索查询"""
    query: str
    max_results: int = 5


# ==================== 文档管理 API ====================

@router.get("/documents")
async def get_documents():
    """
    获取文档列表（返回原始完整文档，而不是分割后的块）
    """
    try:
        kb = get_kb_manager()
        if kb is None:
            raise HTTPException(status_code=503, detail="知识库未初始化")
        
        # 返回原始文档列表
        formatted_docs = []
        for idx, doc in enumerate(kb.original_documents):
            formatted_docs.append({
                "id": doc.get("id", f"doc_{idx}"),
                "content": doc.get("content", ""),
                "source": doc.get("source", "未知"),
                "doc_type": doc.get("doc_type", "未知"),
                "created_at": doc.get("created_at", ""),
                "metadata": {k: v for k, v in doc.get("metadata", {}).items() 
                           if k not in ["source", "doc_type", "created_at"]}
            })
        
        return {
            "data": formatted_docs,
            "total": len(formatted_docs)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取文档列表失败: {str(e)}")


@router.post("/documents")
async def add_documents(batch: DocumentBatchInput):
    """
    添加文档到知识库
    """
    try:
        kb = get_kb_manager()
        if kb is None:
            raise HTTPException(status_code=503, detail="知识库未初始化")
        
        # 转换为字典列表
        documents = [doc.dict() for doc in batch.documents]
        
        # 调用知识库管理器
        result = kb.add_documents(documents)
        
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"添加文档失败: {str(e)}")


@router.delete("/documents/{doc_id}")
async def delete_document(doc_id: str):
    """
    删除单个文档
    """
    try:
        kb = get_kb_manager()
        if kb is None:
            raise HTTPException(status_code=503, detail="知识库未初始化")
        
        # 使用按ID精确删除的方法
        result = kb.delete_documents_by_ids([doc_id])
        
        if result.get("error"):
            raise HTTPException(status_code=400, detail=result.get("error"))
        if result.get("warning"):
            raise HTTPException(status_code=404, detail=result.get("warning"))
        
        return result
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"删除文档失败: {str(e)}")


@router.post("/documents/batch/delete")
async def delete_documents_batch(data: Dict):
    """
    批量删除文档
    """
    try:
        doc_ids = data.get("doc_ids", [])
        kb = get_kb_manager()
        if kb is None:
            raise HTTPException(status_code=503, detail="知识库未初始化")
        
        if not doc_ids:
            return {
                "success": True,
                "deleted_count": 0,
                "message": "没有文档需要删除"
            }
        
        # 使用按ID精确删除的方法
        result = kb.delete_documents_by_ids(doc_ids)
        
        if result.get("error"):
            raise HTTPException(status_code=400, detail=result.get("error"))
        
        return result
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"批量删除失败: {str(e)}")


# ==================== 搜索 API ====================

@router.post("/search")
async def search_documents(query: SearchQuery):
    """
    搜索知识库
    """
    try:
        kb = get_kb_manager()
        if kb is None:
            raise HTTPException(status_code=503, detail="知识库未初始化")
        
        # 调用搜索功能
        results = kb.search_similar(query.query, k=query.max_results, score_threshold=0.3)
        
        return {
            "query": query.query,
            "total_results": len(results),
            "results": results
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"搜索失败: {str(e)}")


# ==================== 统计 API ====================

@router.get("/stats")
async def get_stats():
    """
    获取知识库统计信息
    """
    try:
        kb = get_kb_manager()
        if kb is None:
            raise HTTPException(status_code=503, detail="知识库未初始化")
        
        return kb.get_stats()
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取统计失败: {str(e)}")


# ==================== 导入导出 API ====================

@router.post("/upload")
async def upload_file(
    file: UploadFile = File(...),
    doc_type: Optional[str] = Form(None),
    source: Optional[str] = Form(None),
    metadata: Optional[str] = Form(None)
):
    """
    上传文件
    """
    try:
        logger.info(f"收到文件上传请求: filename={file.filename}, content_type={file.content_type}, size={file.size if hasattr(file, 'size') else 'unknown'}")
        logger.info(f"上传参数: doc_type={doc_type}, source={source}, metadata={metadata}")
        
        # 读取文件内容
        contents = await file.read()
        logger.info(f"文件读取成功，大小: {len(contents)} 字节")
        
        # 使用专门的解析函数处理不同格式的文件
        try:
            content_str = parse_file_content(contents, file.filename, file.content_type)
            logger.info(f"文件解析成功，内容长度: {len(content_str)} 字符")
        except HTTPException:
            # 重新抛出 HTTP 异常
            raise
        except Exception as e:
            error_msg = f"文件解析失败: {str(e)}"
            logger.error(error_msg)
            logger.error(traceback.format_exc())
            raise HTTPException(status_code=400, detail=error_msg)
        
        # 获取知识库管理器
        kb = get_kb_manager()
        if kb is None:
            logger.error("知识库未初始化")
            raise HTTPException(status_code=503, detail="知识库未初始化")
        
        # 解析 metadata
        metadata_dict = {}
        if metadata:
            try:
                metadata_dict = json.loads(metadata)
                logger.info(f"解析 metadata 成功: {metadata_dict}")
            except json.JSONDecodeError as e:
                logger.warning(f"metadata JSON 解析失败: {str(e)}，使用空字典")
                metadata_dict = {}
        
        # 确定文档来源
        doc_source = source or file.filename or "未知来源"
        
        # 确定文档类型
        doc_type_value = doc_type or "other"
        
        # 创建文档
        document = {
            "content": content_str,
            "source": doc_source,
            "doc_type": doc_type_value,
            "metadata": {
                "filename": file.filename,
                "content_type": file.content_type,
                **metadata_dict
            }
        }
        
        logger.info(f"准备添加文档: source={doc_source}, doc_type={doc_type_value}, content_length={len(content_str)}")
        
        # 添加到知识库
        result = kb.add_documents([document])
        
        logger.info(f"文档添加成功: {result}")
        
        return result
        
    except HTTPException:
        # 重新抛出 HTTP 异常
        raise
    except UnicodeDecodeError as e:
        error_msg = f"文件编码错误: {str(e)}"
        logger.error(error_msg)
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=400, detail=error_msg)
    except json.JSONDecodeError as e:
        error_msg = f"metadata JSON 格式错误: {str(e)}"
        logger.error(error_msg)
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=400, detail=error_msg)
    except Exception as e:
        error_msg = f"文件上传失败: {str(e)}"
        logger.error(error_msg)
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=500, detail=error_msg)


@router.post("/import")
async def import_documents(file: UploadFile = File(...)):
    """
    批量导入文档（支持 JSON 格式）
    """
    try:
        contents = await file.read()
        content_str = contents.decode("utf-8")
        
        # 解析 JSON
        documents = json.loads(content_str)
        if not isinstance(documents, list):
            documents = [documents]
        
        kb = get_kb_manager()
        if kb is None:
            raise HTTPException(status_code=503, detail="知识库未初始化")
        
        # 添加到知识库
        result = kb.add_documents(documents)
        
        return result
    except json.JSONDecodeError:
        raise HTTPException(status_code=400, detail="JSON 格式错误")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"导入失败: {str(e)}")


@router.get("/export")
async def export_knowledge_base():
    """
    导出知识库
    """
    try:
        kb = get_kb_manager()
        if kb is None:
            raise HTTPException(status_code=503, detail="知识库未初始化")
        
        # 收集所有文档
        documents = []
        if kb.vectorstore:
            for doc in kb.vectorstore.docstore._dict.values():
                documents.append({
                    "content": doc.page_content,
                    "source": doc.metadata.get("source", ""),
                    "doc_type": doc.metadata.get("doc_type", ""),
                    "metadata": doc.metadata
                })
        
        # 转换为 JSON
        from fastapi.responses import JSONResponse
        return JSONResponse(
            content=documents,
            headers={
                "Content-Disposition": "attachment; filename=knowledge_base.json"
            }
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"导出失败: {str(e)}")


# ==================== 健康检查 ====================

@router.get("/health")
async def health_check():
    """
    知识库健康检查
    """
    try:
        kb = get_kb_manager()
        if kb is None:
            return {
                "status": "unavailable",
                "message": "知识库未初始化"
            }
        
        stats = kb.get_stats()
        return {
            "status": "healthy",
            "message": "知识库正常运行",
            "stats": stats
        }
    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }

